import os
import time
import cv2
import numpy as np
from sklearn import neighbors
from sklearn.svm import SVC
import load_mnist

answer=[0,0,0,0,0,0,0,0,0,0,
        1,1,1,1,1,1,1,1,1,1,
        2,2,2,2,2,2,2,2,2,2,
        3,3,3,3,3,3,3,3,3,3,
        4,4,4,4,4,4,4,4,4,4,
        5,5,5,5,5,5,5,5,5,5,
        6,6,6,6,6,6,6,6,6,6,
        7,7,7,7,7,7,7,7,7,7,
        8,8,8,8,8,8,8,8,8,8,
        9,9,9,9,9,9,9,9,9,9,]
# 图片颜色取反
def colorInversion(img):
    height = img.shape[0]  # 高
    width = img.shape[1]  # 宽
    channels = img.shape[2]  # 通道数
    # 将图像的每个像素点进行反选操作
    for row in range(height):
        for col in range(width):
            for c in range(channels):
                pv = img[row, col, c]
                img[row, col, c] = 255 - pv

    return img

def Knn_prd():
    knn = neighbors.KNeighborsClassifier(algorithm='kd_tree', n_neighbors=3)
    print("数据读取完成，开始训练")
    # print(train_dataSet[:5])
    # print(train_hwLabels[:5])
    T1 = time.time()
    knn.fit(train_dataSet[:train_num], train_hwLabels[:train_num])



    res = knn.predict(list[:test_num])  # 对测试集进行预测
    print(res)
    error_num = np.sum(res != ans[:test_num])  # 统计分类错误的数目

    T2 = time.time()
    print("测试集数量：", test_num, " 错误次数：", error_num, "  准确率：", (test_num - error_num) / float(test_num))
    print("总共用时：", T2 - T1)

def SVM_prd():
    T1 = time.time()
    clf = SVC(C=4, kernel='rbf')
    clf.fit(train_dataSet[:train_num], train_hwLabels[:train_num])


    classifierResult = clf.predict(list[:test_num])
    print(classifierResult)
    error_num = np.sum(classifierResult != ans[:test_num])  # 统计分类错误的数目
    T2 = time.time()
    print("测试集数量：", test_num, " 错误次数：", error_num, "  准确率：", (test_num - error_num) / float(test_num))
    print("总共用时：", T2 - T1)
if __name__ == '__main__':
    path = os.listdir(r'C:\num')

    list = []
    for m in path:
        imagepath = r'C:\num' + '\\' + m

        img = cv2.imread(imagepath)
        img = colorInversion(img)  # 图片颜色取反
        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)  # 把图片转为灰度图  使channels=1

        size = (28, 28)

        shrink = cv2.resize(img, size, interpolation=cv2.INTER_AREA)  # 把图片的大小变成28*28   适应MNIST数据集
        # cv2.imshow("1", shrink)
        # cv2.waitKey()
        a = shrink.reshape(1, 784)  # 把矩阵转成1*784     784=28*28

        list.append(a[0])

    list = np.array(list, dtype=np.uint8)  # 转成uint8的格式
    ans = np.array(answer, dtype=np.uint8)
    print("答案为： ", ans)


    path = r'E:\python\python代码\dataSet'

    train_num = 6000  # 训练集数量     上限为6W
    test_num = 100  # 测试集数量     上限为1W

    # read dataSet
    train_dataSet, train_hwLabels = load_mnist.load_mnist_train(path)
    # read  testing dataSet
    dataSet, hwLabels = load_mnist.load_mnist_test(path)

    Knn_prd()
    SVM_prd()
